package org.loong.crypto.extension.cms;

import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.security.interfaces.ECPublicKey;
import java.util.Collection;
import java.util.Objects;

import org.apache.commons.lang3.ArrayUtils;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.cms.CMSObjectIdentifiers;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSEnvelopedDataGenerator;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.KeyTransRecipientInformation;
import org.bouncycastle.cms.RecipientInformation;
import org.bouncycastle.operator.OutputEncryptor;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.EncryptionMethod;
import org.loong.crypto.core.utils.ECKeyUtils;
import org.loong.crypto.core.utils.KeyTools;
import org.loong.crypto.extension.cms.jce.JceCMSContentEncryptorBuilder;
import org.loong.crypto.extension.cms.jce.JceKeyTransEnvelopedRecipient;
import org.loong.crypto.extension.cms.jce.JceKeyTransRecipient;
import org.loong.crypto.extension.cms.jce.JceKeyTransRecipientInfoGenerator;
import org.loong.crypto.service.CryptoService;
import org.loong.crypto.spi.Join;

/**
 * The envelopedData provider.
 */
@Join
public class EnvelopedDataProvider implements CMSProvider {

    @Override
    public byte[] generate(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待加密的数据不允许为空.");
        }

        if (Objects.isNull(params.getEncryptionMethod())) {
            throw new CryptoException("数据加密算法不允许为空.");
        }

        final X509Certificate encryptionCertificate = params.getEncryptionCertificate();
        final PublicKey publicKey = params.getEncryptionPublicKey();
        if (Objects.isNull(encryptionCertificate) && Objects.isNull(publicKey)) {
            throw new CryptoException("接收者证书或公钥不允许为空.");
        }

        try {
            CMSEnvelopedDataGenerator cmsEnvelopedDataGenerator = new CMSEnvelopedDataGenerator();
            JceKeyTransRecipientInfoGenerator recipientGenerator = null;
            if (Objects.nonNull(encryptionCertificate)) {
                recipientGenerator = new JceKeyTransRecipientInfoGenerator(cryptoService, encryptionCertificate);
            } else if (Objects.nonNull(publicKey)) {
                byte[] keyId = KeyTools.createSubjectKeyId(publicKey).getKeyIdentifier();
                recipientGenerator = new JceKeyTransRecipientInfoGenerator(cryptoService, keyId, publicKey);
            }
            cmsEnvelopedDataGenerator.addRecipientInfoGenerator(recipientGenerator);

            ASN1ObjectIdentifier contentType = CMSObjectIdentifiers.envelopedData;
            if ((Objects.nonNull(encryptionCertificate) && ECKeyUtils.isSM2Algorithm(new ASN1ObjectIdentifier(encryptionCertificate.getSigAlgOID())))
                    || (Objects.nonNull(publicKey) && publicKey instanceof ECPublicKey)) {
                contentType = GMCMSObjectIdentifiers.envelopedData;
            }

            CMSTypedData cmsData = new CMSProcessableByteArray(contentType, params.getData());
            OutputEncryptor encryptor = new JceCMSContentEncryptorBuilder(cryptoService, EncryptionMethod.find(params.getEncryptionMethod())).build();
            CMSEnvelopedData envelopedData = cmsEnvelopedDataGenerator.generate(cmsData, encryptor);
            return envelopedData.getEncoded();
        } catch (CertificateEncodingException | CMSException | IOException | NoSuchAlgorithmException e) {
            e.printStackTrace();
            throw new CryptoException(e.getMessage(), e);
        }
    }

    @SuppressWarnings("unchecked")
    @Override
    public byte[] parse(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待解密的数据不允许为空.");
        }

        final X509Certificate encryptionCertificate = params.getEncryptionCertificate();
        final PublicKey publicKey = params.getEncryptionPublicKey();
        if (Objects.isNull(encryptionCertificate) && Objects.isNull(publicKey)) {
            throw new CryptoException("接收者证书或公钥不允许为空.");
        }

        final PrivateKey privateKey = params.getPrivateKey();
        if (Objects.isNull(privateKey)) {
            throw new CryptoException("接收者私钥不允许为空.");
        }

        try {
            final CMSEnvelopedData envelopedData = new CMSEnvelopedData(params.getData());
            if (Objects.isNull(envelopedData)) {
                throw new CryptoException("数字信封数据内容为空.");
            }

            final Collection<RecipientInformation> recipients = envelopedData.getRecipientInfos().getRecipients();
            if (recipients.isEmpty()) {
                throw new CryptoException("数字信封数据接收者信息为空.");
            }

            X509CertificateHolder encryptionCertificateHolder = null;
            byte[] keyId = null;
            if (Objects.nonNull(encryptionCertificate)) {
                encryptionCertificateHolder = new X509CertificateHolder(encryptionCertificate.getEncoded());
                keyId = KeyTools.createSubjectKeyId(encryptionCertificate.getPublicKey()).getKeyIdentifier();
            } else if (Objects.nonNull(publicKey)) {
                keyId = KeyTools.createSubjectKeyId(publicKey).getKeyIdentifier();
            }

            for (RecipientInformation info : recipients) {
                KeyTransRecipientInformation recipientInfo = (KeyTransRecipientInformation) info;
                if ((Objects.nonNull(encryptionCertificateHolder) && recipientInfo.getRID().match(encryptionCertificateHolder))
                        || (Objects.nonNull(publicKey) && recipientInfo.getRID().match(keyId))) {
                    JceKeyTransRecipient recipient = new JceKeyTransEnvelopedRecipient(cryptoService, privateKey);
                    return recipientInfo.getContent(recipient);
                }
            }
            return null;
        } catch (CMSException | CertificateEncodingException | IOException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }
}
